/*
3d_iou (cpu)
3d_nms (cpu)
fast postprocess(CPU/ARM)
Written by Zhihua Liu
*/

#include <stdio.h>
#include <math.h>
#include <vector>
#include<iostream>
#include <iomanip>
#include <cstring>
#include <algorithm>
#include "cc_decode.h"

#define pi 3.1415926
#define DATASIZE 20
inline float min(float a, float b){
    return a > b ? b : a;
}
//inline float min(const float &a,const float &b){
//    return a > b ? b : a;
//}
inline float max(float a, float b){
    return a > b ? a : b;
}

const float EPS = 1e-8;
struct Point {
    float x, y;
    Point() {}
    Point(double _x, double _y){
        x = _x, y = _y;
    }

     void set(float _x, float _y){
        x = _x; y = _y;
    }

     Point operator +(const Point &b)const{
        return Point(x + b.x, y + b.y);
    }

    Point operator -(const Point &b)const{
        return Point(x - b.x, y - b.y);
    }
};

inline float cross(const Point &a, const Point &b){
    return a.x * b.y - a.y * b.x;
}

inline float cross(const Point &p1, const Point &p2, const Point &p0){
    return (p1.x - p0.x) * (p2.y - p0.y) - (p2.x - p0.x) * (p1.y - p0.y);
}

inline int check_rect_cross(const Point &p1, const Point &p2, const Point &q1, const Point &q2){
    int ret = min(p1.x,p2.x) <= max(q1.x,q2.x) &&
              min(q1.x,q2.x) <= max(p1.x,p2.x) &&
              min(p1.y,p2.y) <= max(q1.y,q2.y) &&
              min(q1.y,q2.y) <= max(p1.y,p2.y);
    return ret;
}

inline int check_in_box2d(const float *box, const Point &p){
    //params: (7) [x, y, z, dx, dy, dz, heading]
    const float MARGIN = 1e-2;

    float center_x = box[0], center_y = box[1];
    float angle_cos = cos(-box[6]), angle_sin = sin(-box[6]);  // rotate the point in the opposite direction of box
    float rot_x = (p.x - center_x) * angle_cos + (p.y - center_y) * (-angle_sin);
    float rot_y = (p.x - center_x) * angle_sin + (p.y - center_y) * angle_cos;

    return (fabs(rot_x) < box[3] / 2 + MARGIN && fabs(rot_y) < box[4] / 2 + MARGIN);
}

inline int intersection(const Point &p1, const Point &p0, const Point &q1, const Point &q0, Point &ans){
    // fast exclusion
    if (check_rect_cross(p0, p1, q0, q1) == 0) return 0;

    // check cross standing
    float s1 = cross(q0, p1, p0);
    float s2 = cross(p1, q1, p0);
    float s3 = cross(p0, q1, q0);
    float s4 = cross(q1, p1, q0);

    if (!(s1 * s2 > 0 && s3 * s4 > 0)) return 0;

    // calculate intersection of two lines
    float s5 = cross(q1, p1, p0);
    if(fabs(s5 - s1) > EPS){
        ans.x = (s5 * q0.x - s1 * q1.x) / (s5 - s1);
        ans.y = (s5 * q0.y - s1 * q1.y) / (s5 - s1);

    }
    else{
        float a0 = p0.y - p1.y, b0 = p1.x - p0.x, c0 = p0.x * p1.y - p1.x * p0.y;
        float a1 = q0.y - q1.y, b1 = q1.x - q0.x, c1 = q0.x * q1.y - q1.x * q0.y;
        float D = a0 * b1 - a1 * b0;

        ans.x = (b0 * c1 - b1 * c0) / D;
        ans.y = (a1 * c0 - a0 * c1) / D;
    }

    return 1;
}

inline void rotate_around_center(const Point &center, const float angle_cos, const float angle_sin, Point &p){
    float new_x = (p.x - center.x) * angle_cos + (p.y - center.y) * (-angle_sin) + center.x;
    float new_y = (p.x - center.x) * angle_sin + (p.y - center.y) * angle_cos + center.y;
    p.set(new_x, new_y);
}

inline int point_cmp(const Point &a, const Point &b, const Point &center){
    return atan2(a.y - center.y, a.x - center.x) > atan2(b.y - center.y, b.x - center.x);
}

inline float L2(const float a, const float b){
    return (a - b) * (a - b);
}
inline float box_overlap(const float *box_a, const float *box_b){
    // params: box_a (7) [x, y, z, dx, dy, dz, heading]
    // params: box_b (7) [x, y, z, dx, dy, dz, heading]

    float dis_center = sqrtf(L2(box_a[0], box_b[0]) + L2(box_a[1], box_b[1]));
    float dis_corner = (box_a[3] + box_a[4] + box_b[3] + box_b[4]) / 2;
    if (dis_center >= dis_corner ) return EPS;

    float a_angle = box_a[6], b_angle = box_b[6];
    float a_dx_half = box_a[3] / 2, b_dx_half = box_b[3] / 2, a_dy_half = box_a[4] / 2, b_dy_half = box_b[4] / 2;
    float a_x1 = box_a[0] - a_dx_half, a_y1 = box_a[1] - a_dy_half;
    float a_x2 = box_a[0] + a_dx_half, a_y2 = box_a[1] + a_dy_half;
    float b_x1 = box_b[0] - b_dx_half, b_y1 = box_b[1] - b_dy_half;
    float b_x2 = box_b[0] + b_dx_half, b_y2 = box_b[1] + b_dy_half;

    Point center_a(box_a[0], box_a[1]);
    Point center_b(box_b[0], box_b[1]);

    Point box_a_corners[5];
    box_a_corners[0].set(a_x1, a_y1);
    box_a_corners[1].set(a_x2, a_y1);
    box_a_corners[2].set(a_x2, a_y2);
    box_a_corners[3].set(a_x1, a_y2);

    Point box_b_corners[5];
    box_b_corners[0].set(b_x1, b_y1);
    box_b_corners[1].set(b_x2, b_y1);
    box_b_corners[2].set(b_x2, b_y2);
    box_b_corners[3].set(b_x1, b_y2);

    // get oriented corners
    float a_angle_cos = cos(a_angle), a_angle_sin = sin(a_angle);
    float b_angle_cos = cos(b_angle), b_angle_sin = sin(b_angle);

    for (int k = 0; k < 4; k++){
        rotate_around_center(center_a, a_angle_cos, a_angle_sin, box_a_corners[k]);
        rotate_around_center(center_b, b_angle_cos, b_angle_sin, box_b_corners[k]);
    }

    box_a_corners[4] = box_a_corners[0];
    box_b_corners[4] = box_b_corners[0];

    // get intersection of lines
    Point cross_points[16];
    Point poly_center;
    int cnt = 0, flag = 0;

    poly_center.set(0, 0);
    for (int i = 0; i < 4; i++){
        for (int j = 0; j < 4; j++){
            flag = intersection(box_a_corners[i + 1], box_a_corners[i], box_b_corners[j + 1], box_b_corners[j], cross_points[cnt]);
            if (flag){
                poly_center = poly_center + cross_points[cnt];
                cnt++;
            }
        }
    }

    // check corners
    for (int k = 0; k < 4; k++){
        if (check_in_box2d(box_a, box_b_corners[k])){
            poly_center = poly_center + box_b_corners[k];
            cross_points[cnt] = box_b_corners[k];
            cnt++;
        }
        if (check_in_box2d(box_b, box_a_corners[k])){
            poly_center = poly_center + box_a_corners[k];
            cross_points[cnt] = box_a_corners[k];
            cnt++;
        }
    }

    poly_center.x /= cnt;
    poly_center.y /= cnt;

    // sort the points of polygon
    Point temp;
    for (int j = 0; j < cnt - 1; j++){
        for (int i = 0; i < cnt - j - 1; i++){
            if (point_cmp(cross_points[i], cross_points[i + 1], poly_center)){
                temp = cross_points[i];
                cross_points[i] = cross_points[i + 1];
                cross_points[i + 1] = temp;
            }
        }
    }

    // get the overlap areas
    float area = 0;
    for (int k = 0; k < cnt - 1; k++){
        area += cross(cross_points[k] - cross_points[0], cross_points[k + 1] - cross_points[0]);
    }

    return fabs(area) / 2.0;
}

inline float iou_bev(const float *box_a, const float *box_b){
    // params: box_a (7) [x, y, z, dx, dy, dz, heading]
    // params: box_b (7) [x, y, z, dx, dy, dz, heading]
    float sa = box_a[3] * box_a[4];
    float sb = box_b[3] * box_b[4];
    float s_overlap = box_overlap(box_a, box_b);
    return s_overlap / fmaxf(sa + sb - s_overlap, EPS);
}

inline float sigmoid(float const x)  { return (1 / (1 + exp(-x))); };

inline float limit_period(float val, float period){
    return val - floorf(val / period) * period;
}

inline float fast_exp(float x)
{
    union {uint32_t i;float f;} v;
    v.i=(1<<23)*(1.4426950409*x+126.93490512f);

    return v.f;
}

int nms_cpu(float const* Box, float const* score, bool* keep, int kInfoNumber,float const& kIouThreshold, int filled_size){

    int *idx = (int *) malloc(sizeof(int) * filled_size);
    std::iota(idx, idx+filled_size, 0);
    std::sort(idx, idx + filled_size,
              [&score](size_t i1, size_t i2) {return score[i1] > score[i2];});

    float iou{0};
    int i1,i2;
    for (std::uint32_t i{0U}; i < filled_size; ++i) {
        i1 = idx[i];
        if (keep[i1]) {
            for (std::uint32_t j{i + 1}; j < filled_size; ++j) {
                i2 = idx[j];
                if (keep[i2]) {
                    iou = iou_bev(&Box[i1 * kInfoNumber], &Box[i2 * kInfoNumber]);
                    if (iou > kIouThreshold) {
                        keep[i2] = false;}
                }
            }
        }
    }
    return 1;
}


int decode(float const* dpuout_cls,float const* dpuout_box, float const* dpuout_dir, float const* dpuout_pos,
    float const* anchors_, float const* ClassThreshold,float const pos_threshold,
    float *scores,float* Bbox,int* labels,int const kMaxResultSize, int filled_size_,
    int const NetLength,int const NetWidth,int const ClassNumber, int kRadianNumber,
           int kInfoNumber,const int downsize,float const* voxels_size,float const* pts_range){
    /*    HWC support
     *  float support
     *  score: H * W * cls * 2
     *  box  : H * W * cls * 2 * 7
     *  dir  : H * W * cls * 2
     *  pos  : H * W * 2
     *  anchor_ ClassNumber*(x,y,z,h)
    */

    int kDirectionNum = 2;
    
    float x_stride = voxels_size[0]  * downsize;
    float y_stride = voxels_size[1]  * downsize;

    float x_offset = x_stride / 2 + pts_range[0];
    float y_offset = y_stride / 2 + pts_range[1];

    float class_threshold[DATASIZE]{0};
    float diagonal[DATASIZE]{0};

    for (std::uint16_t i{0U}; i < ClassNumber; ++i) {
        class_threshold[i] = -1 * log((1 - ClassThreshold[i]) / ClassThreshold[i]);
        diagonal[i] = sqrt(anchors_[i * 4] * anchors_[i * 4] +
                                     anchors_[i * 4 + 1] * anchors_[i * 4 + 1]);

    }


    float pos_threshold_reverse = -1 * log((1 - pos_threshold) / pos_threshold);
    std::uint32_t xy_step_cls{ClassNumber * kRadianNumber};
    std::uint32_t xy_step_box{ClassNumber * kRadianNumber * kInfoNumber};
    std::uint32_t xy_step_dir{ClassNumber * kRadianNumber * kDirectionNum};
    std::uint32_t a_xy_step{ClassNumber * kRadianNumber * kInfoNumber};
    float max_score{0};
    std::int8_t label;
     for (std::uint16_t i{0U}; i < NetLength; ++i) {
        for (std::uint16_t j{0U}; j < NetWidth; ++j) {

//            if (dpuout_pos[(i * NetWidth + j)] <
//              pos_threshold_reverse) {
//            continue;
//          }
          std::uint64_t cls_index_cls{(i * NetWidth + j) * xy_step_cls};
          std::uint64_t cls_index_box{(i * NetWidth + j) * xy_step_box};
          std::uint64_t cls_index_dir{(i * NetWidth + j) * xy_step_dir};
          for (std::uint16_t k{0U}; k < ClassNumber; ++k) {
              for (std::uint16_t l{0U}; l < kRadianNumber; ++l) {
                  // step1: decode score and label
                  std::uint64_t cr_index{cls_index_cls +
                                         (k * kRadianNumber + l)};
                  max_score = dpuout_cls[cr_index];
                  label = k;
                if (max_score < class_threshold[label]) {
                continue;
              }

              // step2: decode box
              if (filled_size_ >= kMaxResultSize) {
                return -1;
              }
              cr_index = cls_index_box + (k * kRadianNumber + l) * kInfoNumber;
              scores[filled_size_] = sigmoid(max_score);
              labels[filled_size_] = label;

              Bbox[filled_size_*kInfoNumber + 0] =
                  (dpuout_box[cr_index + 0] ) * diagonal[k] +
                  j * x_stride + x_offset;
              Bbox[filled_size_*kInfoNumber + 1] =
                  (dpuout_box[cr_index + 1] ) * diagonal[k] +
                  i * y_stride + y_offset;
              Bbox[filled_size_*kInfoNumber + 2] =
                  (dpuout_box[cr_index + 2] ) *
                  anchors_[4 * k + 2] + anchors_[3 + 4 * k] + anchors_[4 * k + 2] / 2;
              Bbox[filled_size_*kInfoNumber + 3] =
                  exp(dpuout_box[cr_index + 3] ) *
                  anchors_[4 * k + 0];
              Bbox[filled_size_*kInfoNumber + 4] =
                  exp(dpuout_box[cr_index + 4] ) *
                          anchors_[4 * k + 1];
              Bbox[filled_size_*kInfoNumber + 5] =
                  exp(dpuout_box[cr_index + 5] ) *
                          anchors_[4 * k + 2 ];
              Bbox[filled_size_*kInfoNumber + 6] =
                      dpuout_box[cr_index + 6] + l * pi / 2 + pi;
//                  limit_period(dpuout_box[cr_index + 6] + l * pi / 2 - pi / 2, pi) + pi/2;

              // step3: decode direction
              cr_index = cls_index_dir +
                  (k * kRadianNumber + l) * kDirectionNum;
              std::int16_t dir{
                  (dpuout_dir[cr_index] > dpuout_dir[cr_index + 1])
                  ? 0.0f : 1.0f};
              Bbox[filled_size_*kInfoNumber + 6] += pi * dir;
              filled_size_++;
            }
          }
        }
     }
     return filled_size_;
}

int decode_chw(float const* dpuout_cls,float const* dpuout_box, float const* dpuout_dir, float const* dpuout_pos,
           float const* anchors_, float const* ClassThreshold,float const pos_threshold,
           float *scores,float* Bbox,int* labels,int const kMaxResultSize, int filled_size_,
           int const NetLength,int const NetWidth, int const ClassNumber, int kRadianNumber,
           int kInfoNumber,const int downsize,float const* voxels_size,float const* pts_range) {
    /*    CHW support
     *  float support
     *  score:  cls * 2     * H * W
     *  box  :  cls * 2 * 7 * H * W
     *  dir  :  cls * 2     * H * W
     *  pos  :  2           * H * W
     *  anchor_ ClassNumber*(x,y,z,h)
    */
    int kDirectionNum = 2;
    
    float x_stride = voxels_size[0] * downsize;
    float y_stride = voxels_size[1] * downsize;

    float x_offset = x_stride / 2 + pts_range[0];
    float y_offset = y_stride / 2 + pts_range[1];

    float class_threshold[DATASIZE]{0};
    float diagonal[DATASIZE]{0};

    for (std::uint16_t i{0U}; i < ClassNumber; ++i) {
        class_threshold[i] = -1 * log((1 - ClassThreshold[i]) / ClassThreshold[i]);
        diagonal[i] = sqrt(anchors_[i * 4] * anchors_[i * 4] +
                           anchors_[i * 4 + 1] * anchors_[i * 4 + 1]);

    }


    float pos_threshold_reverse = -1 * log((1 - pos_threshold) / pos_threshold);

    std::uint32_t xy_step{NetLength * NetWidth};
    std::uint32_t xy_step_box{NetLength * NetWidth * kInfoNumber};
    std::uint32_t xy_step_dir{NetLength * NetWidth * kRadianNumber};

    float max_score{0};
    std::int8_t label;
    for (std::uint16_t i{0U}; i < NetLength; ++i) {
        for (std::uint16_t j{0U}; j < NetWidth; ++j) {
            std::uint64_t index_{(i * NetWidth + j)};
//            if (dpuout_pos[index_] < pos_threshold_reverse) {
//                continue;
//            }
            for (std::uint16_t k{0U}; k < ClassNumber; ++k) {
                for (std::uint16_t l{0U}; l < kRadianNumber; ++l) {
                    // step1: decode score and label
                    std::uint64_t cr_index{index_ +
                                           (k * kRadianNumber + l) * xy_step};
                    max_score = dpuout_cls[cr_index];
                    label = k;

                    if (max_score < class_threshold[label]) {
                        continue;
                    }

                    // step2: decode box
                    if (filled_size_ >= kMaxResultSize) {
                        return -1;
                    }
                    cr_index = index_ + (k * kRadianNumber + l) * xy_step_box;
                    scores[filled_size_] = sigmoid(max_score);
                    labels[filled_size_] = label;

                    Bbox[filled_size_ * kInfoNumber + 0] =
                            (dpuout_box[cr_index + 0 * xy_step]) * diagonal[k] +
                            j * x_stride + x_offset;
                    Bbox[filled_size_ * kInfoNumber + 1] =
                            (dpuout_box[cr_index + 1 * xy_step]) * diagonal[k] +
                            i * y_stride + y_offset;
                    Bbox[filled_size_ * kInfoNumber + 2] =
                            (dpuout_box[cr_index + 2 * xy_step]) *
                            anchors_[4 * k + 2] + anchors_[3 + 4 * k] + anchors_[4 * k + 2] / 2;
                    Bbox[filled_size_ * kInfoNumber + 3] =
                            exp(dpuout_box[cr_index + 3 * xy_step]) *
                            anchors_[4 * k + 0];
                    Bbox[filled_size_ * kInfoNumber + 4] =
                            exp(dpuout_box[cr_index + 4 * xy_step]) *
                            anchors_[4 * k + 1];
                    Bbox[filled_size_ * kInfoNumber + 5] =
                            exp(dpuout_box[cr_index + 5 * xy_step]) *
                            anchors_[4 * k + 2];
                    Bbox[filled_size_ * kInfoNumber + 6] =
                            dpuout_box[cr_index + 6 * xy_step] + l * pi / 2 + pi;
//                  limit_period(dpuout_box[cr_index + 6] + l * pi / 2 - pi / 2, pi) + pi/2;

                    // step3: decode direction
                    cr_index = index_ +
                               (k * kRadianNumber + l) * xy_step_dir;
                    std::int16_t dir{
                            (dpuout_dir[cr_index] > dpuout_dir[cr_index + xy_step])
                            ? 0.0f : 1.0f};
                    Bbox[filled_size_ * kInfoNumber + 6] += pi * dir;
                    filled_size_++;
                }
            }
        }
    }
    return filled_size_;
}
int decode_hwc8(std::int8_t const* dpuout_cls,std::int8_t const* dpuout_box,
                std::int8_t const* dpuout_dir, std::int8_t const* dpuout_pos,
                float const scale_cls,float const scale_box,float const scale_dir,float const scale_pos,
                float const* anchors_, float const* ClassThreshold,float const pos_threshold,
                float *scores,float* Bbox,int* labels,int const kMaxResultSize, int filled_size_,
                int const NetLength,int const NetWidth,int const ClassNumber, int kRadianNumber,
                int kInfoNumber,const int downsize,float const* voxels_size,float const* pts_range){
    /*    HWC support
     *  int8  support
     *  score: H * W * cls * 2
     *  box  : H * W * cls * 2 * 7
     *  dir  : H * W * cls * 2
     *  pos  : H * W * 2
     *  anchor_ ClassNumber*(x,y,z,h)
    */

    int kDirectionNum = 2;
    float x_stride = voxels_size[0]  * downsize;
    float y_stride = voxels_size[1]  * downsize;

    float x_offset = x_stride / 2 + pts_range[0];
    float y_offset = y_stride / 2 + pts_range[1];

    float class_threshold[DATASIZE]{0};
    float diagonal[DATASIZE]{0};

    for (std::uint16_t i{0U}; i < ClassNumber; ++i) {
        class_threshold[i] = -1 * log((1 - ClassThreshold[i]) / ClassThreshold[i]);
        diagonal[i] = sqrt(anchors_[i * 4] * anchors_[i * 4] +
                           anchors_[i * 4 + 1] * anchors_[i * 4 + 1]);

    }


    float pos_threshold_reverse = -1 * log((1 - pos_threshold) / pos_threshold);
    std::uint32_t xy_step_cls{ClassNumber * kRadianNumber};
    std::uint32_t xy_step_box{ClassNumber * kRadianNumber * kInfoNumber};
    std::uint32_t xy_step_dir{ClassNumber * kRadianNumber * kDirectionNum};
    std::uint32_t a_xy_step{ClassNumber * kRadianNumber * kInfoNumber};
    float max_score{0};
    std::int8_t label;
    for (std::uint16_t i{0U}; i < NetLength; ++i) {
        for (std::uint16_t j{0U}; j < NetWidth; ++j) {

//            if (dpuout_pos[(i * NetWidth + j)] <
//              pos_threshold_reverse) {
//            continue;
//          }
            std::uint64_t cls_index_cls{(i * NetWidth + j) * xy_step_cls};
            std::uint64_t cls_index_box{(i * NetWidth + j) * xy_step_box};
            std::uint64_t cls_index_dir{(i * NetWidth + j) * xy_step_dir};
            for (std::uint16_t k{0U}; k < ClassNumber; ++k) {
                for (std::uint16_t l{0U}; l < kRadianNumber; ++l) {
                    // step1: decode score and label
                    std::uint64_t cr_index{cls_index_cls +
                                           (k * kRadianNumber + l)};
                    max_score = dpuout_cls[cr_index] * scale_cls;
                    label = k;
                    if (max_score < class_threshold[label]) {
                        continue;
                    }

                    // step2: decode box
                    // exp / fast_exp
                    if (filled_size_ >= kMaxResultSize) {
                        return -1;
                    }
                    cr_index = cls_index_box + (k * kRadianNumber + l) * kInfoNumber;
                    scores[filled_size_] = sigmoid(max_score);
                    labels[filled_size_] = label;

                    Bbox[filled_size_*kInfoNumber + 0] =
                            (dpuout_box[cr_index + 0] * scale_box ) * diagonal[k] +
                            j * x_stride + x_offset;
                    Bbox[filled_size_*kInfoNumber + 1] =
                            (dpuout_box[cr_index + 1] * scale_box ) * diagonal[k] +
                            i * y_stride + y_offset;
                    Bbox[filled_size_*kInfoNumber + 2] =
                            (dpuout_box[cr_index + 2] * scale_box) *
                            anchors_[4 * k + 2] + anchors_[3 + 4 * k] + anchors_[4 * k + 2] / 2;
                    Bbox[filled_size_*kInfoNumber + 3] =
                            exp(dpuout_box[cr_index + 3] * scale_box ) *
                            anchors_[4 * k + 0];
                    Bbox[filled_size_*kInfoNumber + 4] =
                            exp(dpuout_box[cr_index + 4] * scale_box) *
                            anchors_[4 * k + 1];
                    Bbox[filled_size_*kInfoNumber + 5] =
                            exp(dpuout_box[cr_index + 5] * scale_box) *
                            anchors_[4 * k + 2 ];
                    Bbox[filled_size_*kInfoNumber + 6] =
                            dpuout_box[cr_index + 6] * scale_box + l * pi / 2 + pi;
//                  limit_period(dpuout_box[cr_index + 6] + l * pi / 2 - pi / 2, pi) + pi/2;

                    // step3: decode direction
                    cr_index = cls_index_dir +
                               (k * kRadianNumber + l) * kDirectionNum;
                    std::int16_t dir{
                            (dpuout_dir[cr_index] > dpuout_dir[cr_index + 1])
                            ? 0.0f : 1.0f};
                    Bbox[filled_size_*kInfoNumber + 6] += pi * dir;
                    filled_size_++;
                }
            }
        }
    }
    return filled_size_;
}